# -*- encoding: utf-8 -*-
'''
@File    :   view.py
@Time    :   2021/11/22 9:43
@Author  :   ZhangChaoYang
@Desc    :   
'''

import tensorflow as tf
import matplotlib.pyplot as plt
from sklearn import metrics
import numpy as np
import itertools


def draw_train_history(history, chart_file="", begin_epoch=10):
    history = history.history
    if "loss" in history:
        plt.subplots(1, 1)
        plt.title("LOSS")
        plt.plot(history["loss"][begin_epoch:], color="g", label="train_loss")
        if "val_loss" in history:
            plt.plot(history["val_loss"][begin_epoch:], color="r", label="valid_loss")
        plt.legend(loc="best")
        if chart_file:
            plt.savefig(chart_file, format="png")  # show()要放在savefig后面，否则会出现空白图片

    if "recon" in history:
        plt.subplots(1, 1)
        plt.title("RECONSTRUCT")
        plt.plot(history["recon"][begin_epoch:], color="g", label="train_reconstruct_loss")
        if "val_recon" in history:
            plt.plot(history["val_recon"][begin_epoch:], color="r", label="valid_reconstruct_loss")
        plt.legend(loc="best")
        if chart_file:
            plt.savefig(chart_file + "_recon", format="png")

    if "energy" in history:
        plt.subplots(1, 1)
        plt.title("ENERGY")
        plt.plot(history["energy"][begin_epoch:], color="g", label="train_energy_loss")
        if "val_energy" in history:
            plt.plot(history["val_energy"][begin_epoch:], color="r", label="valid_energy_loss")
        plt.legend(loc="best")
        if chart_file:
            plt.savefig(chart_file + "_energy", format="png")


def draw_fit_error_1d(X, ano_X, X_hat, ano_X_hat, chart_file=""):
    N = min(X.shape[0], ano_X.shape[0], 6)
    X_hat = X_hat[:N]
    ano_X_hat = ano_X_hat[:N]
    fig, axes = plt.subplots(2, N, figsize=(20, 10))
    fig.suptitle(f'拟合情况')
    for i in range(N):
        axes[0][i].set_title(f"正常样本{i}")
        axes[1][i].set_title(f"异常样本{i}")
        if tf.reduce_mean(tf.abs(X[i])).numpy().tolist() > tf.reduce_mean(
                tf.abs(X_hat[i])).numpy().tolist():  # 幅值大的画在图层的下面

            axes[0][i].plot(X[i])
            axes[0][i].plot(X_hat[i])
        else:
            axes[0][i].plot(X_hat[i])
            axes[0][i].plot(X[i])
        if tf.reduce_mean(tf.abs(ano_X[i])).numpy().tolist() > tf.reduce_mean(tf.abs(ano_X_hat[i])).numpy().tolist():
            axes[1][i].plot(ano_X[i])
            axes[1][i].plot(ano_X_hat[i])
        else:
            axes[1][i].plot(ano_X_hat[i])
            axes[1][i].plot(ano_X[i])
    if chart_file:
        plt.savefig(chart_file, format="png")
    # plt.show()


def draw_fit_error_2d(X, ano_X, X_hat, ano_X_hat, chart_file=""):
    N = min(X.shape[0], ano_X.shape[0], 6)
    X_hat = X_hat[:N]
    ano_X_hat = ano_X_hat[:N]
    _, axes = plt.subplots(2, N, figsize=(20, 10))
    for i in range(N // 2):
        axes[0][i].contourf(X[i])
        axes[1][i].contourf(X_hat[i])
        axes[0][i + N // 2].contourf(ano_X[i])
        axes[1][i + N // 2].contourf(ano_X_hat[i])
    if chart_file:
        plt.savefig(chart_file, format="png")
    # plt.show()


def draw_roc(y_true, y_pred, chart_file=""):
    plt.subplots(1, 1)
    fpr, tpr, thresholds = metrics.roc_curve(y_true, y_pred)
    auc = metrics.auc(fpr, tpr)
    plt.title('auc={:.2f}'.format(auc))
    plt.plot(fpr, tpr)
    plt.plot([0, 1], [0, 1], color='navy', linestyle='--')
    plt.xlim([0., 1.])
    plt.ylim([0., 1.05])
    plt.xlabel("False Positive Rate")
    plt.ylabel("True Positive Rate")
    if chart_file:
        plt.savefig(chart_file, format="png")
    return auc


def plot_confusion_matrix(cm, classes,
                          normalize=False,
                          title='Confusion Matrix',
                          cmap=plt.cm.Blues,
                          chart_file=""):
    '''
    画多分类的混淆矩阵
    :param cm:
    :param classes:
    :param normalize:
    :param title:
    :param cmap:
    :param chart_file:
    :return:
    '''
    if normalize:
        cm = (cm.astype('float') / cm.sum(axis=1)[:, np.newaxis])
        print("Normalized confusion matrix")
    else:
        print('Confusion matrix, without normalization')

    print(cm)
    plt.figure(figsize=(12, 9), dpi=80)
    plt.imshow(cm, interpolation='nearest', cmap=cmap)
    plt.title(title)
    plt.colorbar()
    tick_marks = np.arange(len(classes))
    plt.xticks(tick_marks, classes, rotation=45)
    plt.yticks(tick_marks, classes)

    fmt = '0.2f' if normalize else 'd'
    thresh = cm.max() / 2.
    for i, j in itertools.product(range(cm.shape[0]), range(cm.shape[1])):
        plt.text(j, i, format(cm[i, j], fmt),
                 horizontalalignment="center",
                 color="white" if cm[i, j] > thresh else "black")

    plt.tight_layout()
    plt.ylabel('True label')
    plt.xlabel('Predicted label')
    if chart_file:
        plt.savefig(chart_file, format="png")


def draw_fit_error(e1, e2, chart_file):
    plt.subplots(1, 1)
    plt.title("各分位点的拟合误差")
    plt.plot(range(0, 100, 5), e1, color="blue", label="正常样本")
    plt.plot(range(0, 100, 5), e2, color="red", label="故障样本")
    plt.axhline(e1[-1], color='black', linestyle='dashed')
    pct = 0
    for i in range(len(e1)):
        if e2[i] > e1[-1]:
            pct = i * 5
            break
    plt.axvline(pct, color='green', linestyle='dashed', label=f'故障召回率{100-pct}%')
    plt.legend(loc="best")
    if chart_file:
        plt.savefig(chart_file + "_fit_error.png", format="png")
